package com.trendytech.aidata.udf.sql_function;

import com.trendy.bigdata.db.ipv4.maker.IPv4Maker;
import com.trendy.bigdata.db.ipv4.search.IPv4Searcher;
import com.trendy.bigdata.db.ipv6.maker.IPv6Maker;
import com.trendy.bigdata.db.ipv6.search.IPv6Searcher;
import org.apache.flink.table.functions.ScalarFunction;
import sun.net.util.IPAddressUtil;

import java.util.ArrayList;
import java.util.List;
import java.util.Objects;

public class SimplifiedIpSearchFunction extends ScalarFunction {

    // 单个 IPv4 查询引擎
    private IPv4Searcher ipv4Searcher;
    // 单个 IPv6 查询引擎
    private IPv6Searcher ipv6Searcher;

    // 模拟的 IPv4 数据
    private List<String> mockIpv4Data() {
        List<String> ipv4List = new ArrayList<>();
        // 模拟数据格式：start_ip<SEP>end_ip<SEP>other_info
        ipv4List.add("192.168.1.1<SEP>192.168.1.255<SEP>Info1");
        ipv4List.add("10.0.0.1<SEP>10.0.0.255<SEP>Info2");
        return ipv4List;
    }

    // 模拟的 IPv6 数据
    private List<String> mockIpv6Data() {
        List<String> ipv6List = new ArrayList<>();
        // 模拟数据格式：start_ip<SEP>end_ip<SEP>other_info
        ipv6List.add("2001:db8::1<SEP>2001:db8::ff<SEP>Info3");
        ipv6List.add("2001:abcd::1<SEP>2001:abcd::ff<SEP>Info4");
        return ipv6List;
    }

    @Override
    public void open(org.apache.flink.table.functions.FunctionContext context) throws Exception {
        List<String> ipv4List = mockIpv4Data();
        if (Objects.nonNull(ipv4List) && !ipv4List.isEmpty()) {
            // 构建 IPv4 查询引擎
            IPv4Maker ipv4Maker = new IPv4Maker(ipv4List);
            byte[] ipv4Buffer = ipv4Maker.getDynamicByteArray().getBuffer();
            ipv4Searcher = IPv4Searcher.newWithBuffer(ipv4Buffer);
        }

        List<String> ipv6List = mockIpv6Data();
        if (Objects.nonNull(ipv6List) && !ipv6List.isEmpty()) {
            // 构建 IPv6 查询引擎
            IPv6Maker ipv6Maker = new IPv6Maker(ipv6List);
            byte[] ipv6Buffer = ipv6Maker.getDynamicByteArray().getBuffer();
            ipv6Searcher = IPv6Searcher.newWithBuffer(ipv6Buffer);
        }
    }

    public String eval(String ip) {
        if (ip.contains(":")) {
            // 如果是 IPv6 地址
            if (Objects.nonNull(ipv6Searcher)) {
                return ipv6Searcher.search(IPAddressUtil.textToNumericFormatV6(ip));
            }
        } else {
            // 如果是 IPv4 地址
            if (Objects.nonNull(ipv4Searcher)) {
                return ipv4Searcher.search(ip);
            }
        }
        return "";
    }
}